
from tqdm import tqdm
import pandas as pd
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
import seaborn as sns
import matplotlib.ticker as ticker
from global_land_mask import globe
from gsj_typhoon import tydat,see,split_str_id
from matplotlib.colors import BoundaryNorm
from matplotlib.cm import get_cmap
from typlot.config.global_config import *

ini_time_mode = '00_12'
obs_baseline = 'land'
# names = ['2023NH42_36','2023NH47_46'],[42,47]
# tynames,tyids = split_str_id(names)

names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
tynames,tyids = split_str_id(names)

for tyname,tyid in tqdm(zip(tynames,tyids)):
    # 不同起报时间文件路径
    # tyname,tyid = 'dusurui','16'
    dir_dsr = os.path.join(global_ensdir,f'{tyname}_{tyid}')
    name_date = sorted(os.listdir(dir_dsr))
    name_date = [i for i in name_date if i[-2:] in ini_time_mode]
    dir_date= [os.path.join(dir_dsr,f) for f in name_date if f[-2:] in ini_time_mode]
    RIstd = 7  # or RIstd = 15
    pic_savepath = os.path.join(global_picdir,f'{obs_baseline}_heatmap',ini_time_mode,'id_start_RIsum')
    os.makedirs(pic_savepath,exist_ok=True)
        
    # count快速增强
    
    num_array = np.zeros((len(dir_date),52),dtype=int)  #52是成员数目
    num=[]
    for date in dir_date:
        # 得到某一起报时刻的所有集合成员数据
        name_all = [f for f in os.listdir(date) if f.startswith("TRACK")]
        sorted_names = sorted(name_all, key=lambda x: int(x.split('TRACK_ID_')[-1]))
        path_all = [os.path.join(date, f) for f in sorted_names ]
        
        num2=[]
        #读取并处理单个成员，借助类来封装处理方式
        for path in path_all:
            t=tydat(path,RIstd)
            n=np.sum(t.num_rapidgrow())
            nn = int(path.split("TRACK_ID_")[-1] )
            num2.append([n,nn]) # n代表某一成员的最大增强次数,nn代表该成员序号
            
        num.append(num2)
    
    
    
    for i in range( len(dir_date) ):
        t = num[i]
        for j in range( len(t) ):
            index = num[i][j][1]
            number = num[i][j][0]
            num_array[i,index] = number
    
    
    
    
    # heatmap  x--Trackid  y-date
    plt.figure(figsize=(20, 15))
    ax = sns.heatmap(num_array, annot=False, fmt="d", cmap='YlGnBu', 
                     cbar_kws={'label': 'number of rapid intensification'}, 
                     linewidths=0.3,annot_kws={"fontsize": 5},
                     vmin=0,vmax=5) # vmin=0,vmax=5
    
    ax.set_xlabel("TRACK ID", fontsize=12)
    ax.set_ylabel("Start time", fontsize=12)
    ax.set_title(f"{tyname} RI{RIstd}  RIpoints in all members      forecast_ini_time:{ini_time_mode}",fontsize=20)
    # y坐标作为时间轴，先处理好datetime对象
    name_date_datetime = [ datetime.strptime(f, "%Y%m%d%H") for f in name_date ]
    ax.set_yticklabels(name_date_datetime, rotation=0) 
    cbar = ax.collections[0].colorbar
    # x axis
    ax.set_xticks(np.arange(num_array.shape[1])) 
    ax.set_xticklabels(np.arange(num_array.shape[1]))
    # 设置 colorbar 的刻度为整数
    cbar.locator = ticker.MaxNLocator(integer=True)
    cbar.update_ticks()
    # plt.show()
    plt.savefig(f"{pic_savepath}\\{tyname}.png",dpi=900)
    plt.close()
    
